from __future__ import print_function
import torch
import numpy as np
from PIL import Image
import os
import cv2
# import pclpy
# from pclpy.visualization import PCLVisualizer

def save_images(save_dir, visuals, image_name, image_size, prob_map):
    """save images to disk"""
    image_name = image_name[0]
    oriSize = (image_size[0].item(), image_size[1].item())
    palet_file = 'datasets/palette.txt'
    impalette = list(np.genfromtxt(palet_file, dtype=np.uint8).reshape(3*256))

    for label, im_data in visuals.items():
        if label == 'output':
            if prob_map:
                im = tensor2confidencemap(im_data)
                im = cv2.resize(im, oriSize)
                cv2.imwrite(os.path.join(save_dir, image_name[:-10]+'road_'+image_name[-10:]), im)
            else:
                im = tensor2labelim(im_data, impalette)
                im = cv2.resize(im, oriSize)
                cv2.imwrite(os.path.join(save_dir, image_name), cv2.cvtColor(im, cv2.COLOR_RGB2BGR))

def tensor2im(input_image, imtype=np.uint8):
    """Converts a image Tensor into an image array (numpy)"""
    if isinstance(input_image, torch.Tensor):
        image_tensor = input_image.data
    else:
        return input_image
    image_numpy = image_tensor[0].cpu().float().numpy()
    if image_numpy.shape[0] == 1:
        image_numpy = np.tile(image_numpy, (3, 1, 1))
    image_numpy = (np.transpose(image_numpy, (1, 2, 0)))* 255.0
    return image_numpy.astype(imtype)

def tensor2labelim(label_tensor, impalette, imtype=np.uint8):
    """Converts a label Tensor into an image array (numpy),
    we use a palette to color the label images"""
    if len(label_tensor.shape) == 4:
        _, label_tensor = torch.max(label_tensor.data.cpu(), 1)

    label_numpy = label_tensor[0].cpu().float().detach().numpy()
    label_image = Image.fromarray(label_numpy.astype(np.uint8))
    label_image = label_image.convert("P")
    label_image.putpalette(impalette)
    label_image = label_image.convert("RGB")
    return np.array(label_image).astype(imtype)

def tensor2confidencemap(label_tensor, imtype=np.uint8):
    """Converts a prediction Tensor into an image array (numpy),
    we output predicted probability maps for kitti submission"""
    softmax_numpy = label_tensor[0].cpu().float().detach().numpy()
    softmax_numpy = np.exp(softmax_numpy)
    label_image = np.true_divide(softmax_numpy[1], softmax_numpy[0] + softmax_numpy[1])
    label_image = np.floor(255 * (label_image - label_image.min()) / (label_image.max() - label_image.min()))
    return np.array(label_image).astype(imtype)


def print_current_losses(epoch, i, losses, t, t_data):
    message = '(epoch: %d, iters: %d, time: %.3f, data: %.3f) ' % (epoch, i, t, t_data)
    for k, v in losses.items():
        message += '%s: %.3f ' % (k, v)
    print(message)


def mkdirs(paths):
    if isinstance(paths, list) and not isinstance(paths, str):
        for path in paths:
            mkdir(path)
    else:
        mkdir(paths)

def mkdir(path):
    if not os.path.exists(path):
        os.makedirs(path)


def confusion_matrix(x, y, n, ignore_label=None, mask=None):
    if mask is None:
        mask = np.ones_like(x) == 1
    k = (x >= 0) & (y < n) & (x != ignore_label) & (mask.astype(np.bool))
    return np.bincount(n * x[k].astype(int) + y[k], minlength=n**2).reshape(n, n)

def getScores(conf_matrix):
    if conf_matrix.sum() == 0:
        return 0, 0, 0, 0, 0
    with np.errstate(divide='ignore',invalid='ignore'):
        globalacc = np.diag(conf_matrix).sum() / np.float(conf_matrix.sum())
        classpre = np.diag(conf_matrix) / conf_matrix.sum(0).astype(np.float)
        classrecall = np.diag(conf_matrix) / conf_matrix.sum(1).astype(np.float)
        IU = np.diag(conf_matrix) / (conf_matrix.sum(1) + conf_matrix.sum(0) - np.diag(conf_matrix)).astype(np.float)
        pre = classpre[1]
        recall = classrecall[1]
        iou = IU[1]
        F_score = 2*(recall*pre)/(recall+pre)
    return globalacc, pre, recall, F_score, iou

'''
一些可视化
'''
import open3d as o3d

def cloudview(points):
    # # 创建一个3D图形对象
    # fig = plt.figure()
    # ax = fig.add_subplot(111, projection='3d')
    # # 在图形中添加点云数据
    # ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=0.1)
    # # 设置坐标轴标签
    # ax.set_xlabel('X Label')
    # ax.set_ylabel('Y Label')
    # ax.set_zlabel('Z Label')
    # # 显示图形
    # plt.show()
    num_points = len(points)
    colors = np.random.rand(num_points, 3)
    # 创建Open3D点云对象
    pcd = o3d.geometry.PointCloud()
    # 将随机点云数据和颜色添加到Open3D点云对象中
    pcd.points = o3d.utility.Vector3dVector(points[:, :3])
    # pcd.colors = o3d.utility.Vector3dVector(colors)
    # 可视化彩色点云
    o3d.visualization.draw_geometries([pcd])

def ndarrayview(label, transpose=False, det_label=None):
    '''
    Args:
        label:
        transpose: 是否行列交换
    Returns:
    '''
    import matplotlib.pyplot as plt
    import matplotlib.cm as mcm
    if(len(label.shape)==3):
        arrs = label.copy()
        # 获取数组的长度，以便设置子图的行数
        n = len(arrs)
        # 计算子图网格的行数和列数
        rows = n // 2 + n % 2
        cols = 2 if n > 1 else 1
        # 创建带有子图网格的图形
        fig, axes = plt.subplots(nrows=rows, ncols=cols,figsize=(20, 10))
        # 在每个子图上绘制一个 ndarray
        for i in range(n):
            row = i // cols
            col = i % cols
            ax = axes[row, col]
            arr = arrs[i]
            if(transpose):
                arr = arr.T
            ax.imshow(arr)
            ax.set_title(f"Array {i + 1}")
        plt.tight_layout()
        plt.show()
    else:
        # 绘制热力图
        label_new = np.expand_dims(label,axis=2)
        plt.imshow(label_new,mcm.viridis)
        plt.colorbar()
        plt.show()

def ndarrayrgbview(label, transpose:bool=False, det_label=None):
    import matplotlib.pyplot as plt
    import matplotlib.cm as mcm
    # 绘制热力图
    assert len(label.shape)==3 and label.shape[2]==3,"array形状应该是[n,m,3],实际是{}".format(label.shape)
    if transpose:
        label=np.transpose(label, (1, 0, 2))
    if det_label is not None:
        fig, ax = plt.subplots()
        # 绘制框  注意array中xy方向，和rect中xy方向，plot中行为array的x方向和rect的y方向
        for box in det_label:
            rect = plt.Rectangle((int(box[1]), int(box[2])), int(box[3]) - int(box[1]), int(box[4]) - int(box[2]),
                                 edgecolor='r', facecolor='none')
            ax.add_patch(rect)
    plt.imshow(label)
    plt.colorbar()
    plt.show()

def arrayview(label, det_label, transpose:bool=False):
    import matplotlib.pyplot as plt
    # 绘制热力图
    assert len(label.shape)==2
    label2 = np.copy(label)
    fig, ax = plt.subplots()
    # 绘制框  注意array中xy方向，和rect中xy方向，plot中行为array的x方向和rect的y方向
    for box in det_label:
        rect = plt.Rectangle((int(box[1]), int(box[2])), int(box[3])-int(box[1]), int(box[4])-int(box[2]), edgecolor='r', facecolor='none')
        ax.add_patch(rect)
    if transpose:
        label2=label.T
    plt.imshow(label2)
    plt.colorbar()
    plt.show()

def view_depth(depth):
    depth_image = o3d.geometry.Image(depth.astype(np.float32))
    # Convert the depth image to point cloud
    intrinsic = o3d.camera.PinholeCameraIntrinsic(
        o3d.camera.PinholeCameraIntrinsicParameters.PrimeSenseDefault)
    point_cloud = o3d.geometry.PointCloud.create_from_depth_image(
        depth_image, intrinsic)
    o3d.visualization.draw_geometries([point_cloud])

# def cloudAndNormal_view(cloud_data,normal_data):
#     # 将数据转换为 pcl.PointNormal 类型
#     point_normals = np.zeros(cloud_data.shape[0], dtype=pclpy.pcl.PointNormal)
#     for i in range(len(cloud_data)):
#         point_normals[i].x, point_normals[i].y, point_normals[i].z = cloud_data[i]
#         point_normals[i].normal_x, point_normals[i].normal_y, point_normals[i].normal_z = normal_data[i]
#     # 创建 pcl.PointCloud 对象
#     cloud = pclpy.pcl.PointCloud.PointNormal()
#     cloud.from_array(point_normals)
#     # 创建可视化对象
#     viewer = PCLVisualizer('Viewer')
#     # 添加点云和法向量显示到窗口
#     viewer.addPointCloudNormals(cloud, 1, 0.05, 'cloud')
#     # 启动可视化窗口
#     viewer.spin()
